// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package org.apache.doris.nereids.trees.expressions;

import org.apache.doris.common.Config;
import org.apache.doris.nereids.analyzer.PlaceholderExpression;
import org.apache.doris.nereids.analyzer.Unbound;
import org.apache.doris.nereids.analyzer.UnboundVariable;
import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.trees.AbstractTreeNode;
import org.apache.doris.nereids.trees.expressions.ArrayItemReference.ArrayItemSlot;
import org.apache.doris.nereids.trees.expressions.functions.ExpressionTrait;
import org.apache.doris.nereids.trees.expressions.functions.Nondeterministic;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
import org.apache.doris.nereids.trees.expressions.functions.scalar.Lambda;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
import org.apache.doris.nereids.trees.expressions.shape.LeafExpression;
import org.apache.doris.nereids.trees.expressions.typecoercion.ExpectsInputTypes;
import org.apache.doris.nereids.trees.expressions.typecoercion.TypeCheckResult;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
import org.apache.doris.nereids.trees.plans.commands.info.PartitionDefinition.MaxValue;
import org.apache.doris.nereids.types.ArrayType;
import org.apache.doris.nereids.types.DataType;
import org.apache.doris.nereids.types.MapType;
import org.apache.doris.nereids.types.StructField;
import org.apache.doris.nereids.types.StructType;
import org.apache.doris.nereids.util.ExpressionUtils;
import org.apache.doris.nereids.util.Utils;

import com.google.common.base.Preconditions;
import com.google.common.base.Suppliers;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.ImmutableSet.Builder;
import com.google.common.collect.Lists;
import org.apache.commons.lang3.StringUtils;

import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.function.Supplier;

/**
 * Abstract class for all Expression in Nereids.
 */
public abstract class Expression extends AbstractTreeNode<Expression> implements ExpressionTrait {
    public static final String DEFAULT_EXPRESSION_NAME = "expression";
    // Mask this expression is generated by rule, should be removed.
    protected Optional<String> exprName = Optional.empty();
    private final int depth;
    private final int width;
    // Mark this expression is from predicate infer or something else infer
    private final boolean inferred;
    private final boolean hasUnbound;
    private final boolean compareWidthAndDepth;
    private final Supplier<Set<Slot>> inputSlots = Suppliers.memoize(
            () -> collect(e -> e instanceof Slot && !(e instanceof ArrayItemSlot)));
    private final int fastChildrenHashCode;

    protected Expression(Expression... children) {
        super(children);

        boolean hasUnbound = false;
        switch (children.length) {
            case 0:
                this.depth = 1;
                this.width = 1;
                this.compareWidthAndDepth = supportCompareWidthAndDepth();
                this.fastChildrenHashCode = 0;
                break;
            case 1:
                Expression child = children[0];
                this.depth = child.depth + 1;
                this.width = child.width;
                this.compareWidthAndDepth = child.compareWidthAndDepth && supportCompareWidthAndDepth();
                this.fastChildrenHashCode = child.fastChildrenHashCode() + 1;
                break;
            case 2:
                Expression left = children[0];
                Expression right = children[1];
                this.depth = Math.max(left.depth, right.depth) + 1;
                this.width = left.width + right.width;
                this.compareWidthAndDepth =
                        left.compareWidthAndDepth && right.compareWidthAndDepth && supportCompareWidthAndDepth();
                this.fastChildrenHashCode = left.fastChildrenHashCode() + right.fastChildrenHashCode() + 2;
                break;
            default:
                int maxChildDepth = 0;
                int sumChildWidth = 0;
                boolean compareWidthAndDepth = true;
                int fastChildrenHashCode = 0;
                for (Expression expression : children) {
                    child = expression;
                    maxChildDepth = Math.max(child.depth, maxChildDepth);
                    sumChildWidth += child.width;
                    hasUnbound |= child.hasUnbound;
                    compareWidthAndDepth &= child.compareWidthAndDepth;
                    fastChildrenHashCode = fastChildrenHashCode + expression.fastChildrenHashCode() + 1;
                }
                this.depth = maxChildDepth + 1;
                this.width = sumChildWidth;
                this.compareWidthAndDepth = compareWidthAndDepth;
                this.fastChildrenHashCode = fastChildrenHashCode;
        }

        checkLimit();
        this.inferred = false;
        this.hasUnbound = hasUnbound || this instanceof Unbound;
    }

    protected Expression(List<Expression> children) {
        this(children, false);
    }

    protected Expression(List<Expression> children, boolean inferred) {
        super(children);

        boolean hasUnbound = false;
        switch (children.size()) {
            case 0:
                this.depth = 1;
                this.width = 1;
                this.compareWidthAndDepth = supportCompareWidthAndDepth();
                this.fastChildrenHashCode = 0;
                break;
            case 1:
                Expression child = children.get(0);
                this.depth = child.depth + 1;
                this.width = child.width;
                this.compareWidthAndDepth = child.compareWidthAndDepth && supportCompareWidthAndDepth();
                this.fastChildrenHashCode = child.fastChildrenHashCode() + 1;
                break;
            case 2:
                Expression left = children.get(0);
                Expression right = children.get(1);
                this.depth = Math.max(left.depth, right.depth) + 1;
                this.width = left.width + right.width;
                this.compareWidthAndDepth =
                        left.compareWidthAndDepth && right.compareWidthAndDepth && supportCompareWidthAndDepth();
                this.fastChildrenHashCode = left.fastChildrenHashCode() + right.fastChildrenHashCode() + 2;
                break;
            default:
                int maxChildDepth = 0;
                int sumChildWidth = 0;
                boolean compareWidthAndDepth = true;
                int fastChildrenhashCode = 0;
                for (Expression expression : children) {
                    child = expression;
                    maxChildDepth = Math.max(child.depth, maxChildDepth);
                    sumChildWidth += child.width;
                    hasUnbound |= child.hasUnbound;
                    compareWidthAndDepth &= child.compareWidthAndDepth;
                    fastChildrenhashCode = fastChildrenhashCode + expression.fastChildrenHashCode() + 1;
                }
                this.depth = maxChildDepth + 1;
                this.width = sumChildWidth;
                this.compareWidthAndDepth = compareWidthAndDepth && supportCompareWidthAndDepth();
                this.fastChildrenHashCode = fastChildrenhashCode;
        }

        checkLimit();
        this.inferred = inferred;
        this.hasUnbound = hasUnbound || this instanceof Unbound;
    }

    private void checkLimit() {
        if (depth > Config.expr_depth_limit) {
            throw new AnalysisException(String.format("Exceeded the maximum depth of an "
                    + "expression tree (%s).", Config.expr_depth_limit));
        }
        if (width > Config.expr_children_limit) {
            throw new AnalysisException(String.format("Exceeded the maximum children of an "
                    + "expression tree (%s).", Config.expr_children_limit));
        }
    }

    public Alias alias(String alias) {
        return new Alias(this, alias);
    }

    // Name of expr, this is used by generating column name automatically when there is no
    // alias
    public String getExpressionName() {
        if (!this.exprName.isPresent()) {
            this.exprName = Optional.of(Utils.normalizeName(this.getClass().getSimpleName(), DEFAULT_EXPRESSION_NAME));
        }
        return this.exprName.get();
    }

    /**
     * check input data types
     */
    public TypeCheckResult checkInputDataTypes() {
        // check all of its children recursively.
        for (Expression child : this.children) {
            TypeCheckResult childResult = child.checkInputDataTypes();
            if (childResult.failed()) {
                return childResult;
            }
        }
        if (this instanceof ExpectsInputTypes) {
            ExpectsInputTypes expectsInputTypes = (ExpectsInputTypes) this;
            TypeCheckResult commonCheckResult = checkInputDataTypesWithExpectTypes(
                    children, expectsInputTypes.expectedInputTypes());
            if (commonCheckResult.failed()) {
                return commonCheckResult;
            }
        }
        return checkInputDataTypesInternal();
    }

    public int fastChildrenHashCode() {
        return fastChildrenHashCode;
    }

    protected TypeCheckResult checkInputDataTypesInternal() {
        return TypeCheckResult.SUCCESS;
    }

    private boolean checkInputDataTypesWithExpectType(DataType input, DataType expected) {
        if (input instanceof ArrayType && expected instanceof ArrayType) {
            return checkInputDataTypesWithExpectType(
                    ((ArrayType) input).getItemType(), ((ArrayType) expected).getItemType());
        } else if (input instanceof MapType && expected instanceof MapType) {
            return checkInputDataTypesWithExpectType(
                    ((MapType) input).getKeyType(), ((MapType) expected).getKeyType())
                    && checkInputDataTypesWithExpectType(
                    ((MapType) input).getValueType(), ((MapType) expected).getValueType());
        } else if (input instanceof StructType && expected instanceof StructType) {
            List<StructField> inputFields = ((StructType) input).getFields();
            List<StructField> expectedFields = ((StructType) expected).getFields();
            if (inputFields.size() != expectedFields.size()) {
                return false;
            }
            for (int i = 0; i < inputFields.size(); i++) {
                if (!checkInputDataTypesWithExpectType(
                        inputFields.get(i).getDataType(),
                        expectedFields.get(i).getDataType())) {
                    return false;
                }
            }
            return true;
        } else {
            return checkPrimitiveInputDataTypesWithExpectType(input, expected);
        }
    }

    private boolean checkPrimitiveInputDataTypesWithExpectType(DataType input, DataType expected) {
        // support fast check the case: input=TinyIntType, expected=NumericType, for example: `1 + 1`.
        // if no this check, there will have an exception when invoke NumericType.toCatalogDataType,
        // when there has lots of expression, the exception become the bottleneck, because an exception
        // need to record the whole StackFrame.
        if (expected.acceptsType(input)) {
            return true;
        }

        // TODO: complete the cast logic like FunctionCallExpr.analyzeImpl
        try {
            return input.toCatalogDataType().matchesType(expected.toCatalogDataType());
        } catch (Throwable t) {
            return false;
        }
    }

    private TypeCheckResult checkInputDataTypesWithExpectTypes(
            List<Expression> inputs, List<DataType> expectedTypes) {
        Preconditions.checkArgument(inputs.size() == expectedTypes.size());
        List<String> errorMessages = Lists.newArrayList();
        for (int i = 0; i < inputs.size(); i++) {
            Expression input = inputs.get(i);
            DataType expected = expectedTypes.get(i);
            if (!checkInputDataTypesWithExpectType(input.getDataType(), expected)) {
                errorMessages.add(String.format("argument %d requires %s type, however '%s' is of %s type",
                        i + 1, expected.simpleString(), input.toSql(), input.getDataType().simpleString()));
            }
        }
        if (!errorMessages.isEmpty()) {
            return new TypeCheckResult(false, StringUtils.join(errorMessages, ", "));
        }
        return TypeCheckResult.SUCCESS;
    }

    public abstract <R, C> R accept(ExpressionVisitor<R, C> visitor, C context);

    @Override
    public List<Expression> children() {
        return children;
    }

    @Override
    public Expression child(int index) {
        return children.get(index);
    }

    public int getWidth() {
        return width;
    }

    public int getDepth() {
        return depth;
    }

    public boolean isInferred() {
        return inferred;
    }

    @Override
    public Expression withChildren(List<Expression> children) {
        throw new RuntimeException();
    }

    public Expression withInferred(boolean inferred) {
        throw new RuntimeException("current expression has not impl the withInferred method");
    }

    /**
     * Whether the expression is a constant.
     */
    public boolean isConstant() {
        if (this instanceof AggregateFunction
                || this instanceof Lambda
                || this instanceof MaxValue
                || this instanceof OrderExpression
                || this instanceof PlaceholderExpression
                || this instanceof Properties
                || this instanceof SubqueryExpr
                || this instanceof UnboundVariable
                || this instanceof Variable
                || this instanceof VariableDesc
                || this instanceof WindowExpression
                || this instanceof WindowFrame) {
            // agg_fun(literal) is not constant, the result depends on the group by keys
            return false;
        }
        if (this instanceof LeafExpression) {
            return this instanceof Literal;
        } else {
            return !(this instanceof Nondeterministic) && ExpressionUtils.allMatch(children(), Expression::isConstant);
        }
    }

    public boolean isZeroLiteral() {
        return this instanceof Literal && ((Literal) this).isZero();
    }

    public final Expression castTo(DataType targetType) throws AnalysisException {
        return uncheckedCastTo(targetType);
    }

    public Expression checkedCastTo(DataType targetType) throws AnalysisException {
        return castTo(targetType);
    }

    protected Expression uncheckedCastTo(DataType targetType) throws AnalysisException {
        throw new RuntimeException("Do not implement uncheckedCastTo");
    }

    /**
     * Get all the input slots of the expression.
     * <p>
     * Note that the input slots of subquery's inner plan is not included.
     */
    public final Set<Slot> getInputSlots() {
        return inputSlots.get();
    }

    /**
     * Get all the input slot ids of the expression.
     * <p>
     * Note that the input slots of subquery's inner plan is not included.
     */
    public final Set<ExprId> getInputSlotExprIds() {
        Set<Slot> inputSlots = getInputSlots();
        Builder<ExprId> exprIds = ImmutableSet.builderWithExpectedSize(inputSlots.size());
        for (Slot inputSlot : inputSlots) {
            exprIds.add(inputSlot.getExprId());
        }
        return exprIds.build();
    }

    public boolean isLiteral() {
        return this instanceof Literal;
    }

    public boolean isNullLiteral() {
        return this instanceof NullLiteral;
    }

    public boolean isSlot() {
        return this instanceof Slot;
    }

    public boolean isColumnFromTable() {
        return (this instanceof SlotReference) && ((SlotReference) this).getColumn().isPresent();
    }

    public boolean isKeyColumnFromTable() {
        return (this instanceof SlotReference) && ((SlotReference) this).getColumn().isPresent()
                && ((SlotReference) this).getColumn().get().isKey();
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }
        Expression that = (Expression) o;
        if ((compareWidthAndDepth
                && (this.width != that.width || this.depth != that.depth
                    || this.fastChildrenHashCode != that.fastChildrenHashCode))
                || arity() != that.arity() || !extraEquals(that)) {
            return false;
        }
        return equalsChildren(that);
    }

    protected boolean equalsChildren(Expression that) {
        List<Expression> children = children();
        List<Expression> thatChildren = that.children();
        for (int i = 0; i < children.size(); i++) {
            if (!children.get(i).equals(thatChildren.get(i))) {
                return false;
            }
        }
        return true;
    }

    protected boolean extraEquals(Expression that) {
        return true;
    }

    @Override
    public int hashCode() {
        return getClass().hashCode() + fastChildrenHashCode();
    }

    /**
     * This expression has unbound symbols or not.
     */
    public boolean hasUnbound() {
        return this.hasUnbound;
    }

    public String shapeInfo() {
        return toSql();
    }

    protected boolean supportCompareWidthAndDepth() {
        return true;
    }
}
